import pandas as pd
import numpy as np
import os

# --- 1. 配置参数 ---

# 定义要生成的股票列表（符合A股代码格式）
STOCKS_TO_GENERATE = [
    '600036.SH', '601318.SH', '000651.SZ', '600519.SH',
    '600000.SH', '600887.SH', '000333.SZ', '002415.SZ',
    '600276.SH', '000002.SZ', '601668.SH', '000725.SZ'
]

# 定义数据的时间范围
START_DATE = '2019-01-01'
END_DATE = '2024-12-31'

# 定义输出文件夹
OUTPUT_DIR = 'data'


def generate_single_stock_data(stock_code: str, start: str, end: str) -> pd.DataFrame:
    """
    为单只股票生成模拟的OHLCV数据。

    Args:
        stock_code (str): 股票代码.
        start (str): 开始日期 'YYYY-MM-DD'.
        end (str): 结束日期 'YYYY-MM-DD'.

    Returns:
        pd.DataFrame: 包含标准格式的股票数据.
    """
    # 创建业务日期范围
    dates = pd.date_range(start=start, end=end, freq='B')
    n_days = len(dates)

    # --- 模拟价格序列 (核心部分) ---
    # 为每只股票随机设定不同的基础参数，使其看起来不同
    initial_price = np.random.uniform(20, 200)
    annual_drift = np.random.uniform(-0.1, 0.4)  # 年化漂移率 (趋势)
    annual_volatility = np.random.uniform(0.15, 0.50) # 年化波动率

    # 转换为日度参数
    daily_drift = annual_drift / 252
    daily_volatility = annual_volatility / np.sqrt(252)

    # 创建一个长周期的正弦波来模拟牛熊市
    cycle_length = np.random.uniform(300, 800)
    long_term_cycle = 0.001 * np.sin(np.arange(n_days) * 2 * np.pi / cycle_length)

    # 生成随机冲击
    random_shocks = np.random.randn(n_days) * daily_volatility
    
    # 合成每日收益率并计算收盘价
    daily_returns = daily_drift + long_term_cycle + random_shocks
    close_prices = initial_price * (1 + daily_returns).cumprod()

    df = pd.DataFrame(index=dates, data={'close': close_prices})

    # --- 基于收盘价生成 OHL 数据 ---
    # 开盘价 = 上一个收盘价 + 小幅随机波动
    open_noise = 1 + np.random.randn(n_days) * 0.005
    df['open'] = df['close'].shift(1) * open_noise
    df.iloc[0, df.columns.get_loc('open')] = initial_price * open_noise[0] # 处理第一个开盘价

    # 最高/最低价基于开盘和收盘价生成
    high_noise = 1 + np.random.uniform(0, 0.015, size=n_days)
    low_noise = 1 - np.random.uniform(0, 0.015, size=n_days)
    df['high'] = df[['open', 'close']].max(axis=1) * high_noise
    df['low'] = df[['open', 'close']].min(axis=1) * low_noise

    # --- 模拟成交量 ---
    base_volume = np.random.randint(1_000_000, 10_000_000)
    volume_noise = base_volume * (0.5 + np.random.rand(n_days))
    # 成交量在价格波动大时会放大
    volume_spikes = 1 + abs(df['close'].pct_change().fillna(0)) * np.random.uniform(5, 15)
    df['volume'] = (volume_noise * volume_spikes).astype(int)

    # --- 格式化输出 ---
    df.reset_index(inplace=True)
    df.rename(columns={'index': 'date'}, inplace=True)

    # 确保列的顺序正确，并进行数据取整
    final_cols = ['date', 'open', 'high', 'low', 'close', 'volume']
    df = df[final_cols]
    df[['open', 'high', 'low', 'close']] = df[['open', 'high', 'low', 'close']].round(2)

    return df


def generate_all_stock_data():
    """
    主函数，生成所有配置好的股票数据并保存到CSV文件。
    """
    # 确保输出目录存在
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    print(f"数据将被保存在 '{OUTPUT_DIR}/' 文件夹中。")

    for stock in STOCKS_TO_GENERATE:
        print(f"正在生成 {stock} 的数据...")
        try:
            stock_df = generate_single_stock_data(stock, START_DATE, END_DATE)
            file_path = os.path.join(OUTPUT_DIR, f"{stock}.csv")
            stock_df.to_csv(file_path, index=False)
            print(f"  -> 已成功保存到 {file_path}")
        except Exception as e:
            print(f"  -> 生成 {stock} 数据时发生错误: {e}")

    print("\n所有示例数据生成完毕！")


if __name__ == "__main__":
    generate_all_stock_data()